package cn.lishiyuan.algorithm.tree;


import java.util.*;
import java.util.function.Function;

public class TrieTree {
    private final Node root;
    // 分词
    private final Function<String,List<String>> split;
    // 合并词
    private final Function<List<String>,String> merge;

    public TrieTree(Function<String,List<String>> split,Function<List<String>,String> merge){
        this.split = split;
        this.merge = merge;
        root = new Node(null);
    }

    private static class Node{
        String word;
        // 子节点
        Map<String, Node> children;

        // 是否是最后一个字母或者单词
        boolean isEndWord = false;

        Node(String word){
            this.word = word;
            this.children = new HashMap<>();
        }
    }


    public void insert(String words) {

        if (words == null || words.trim().isEmpty()){
            return;
        }

        // 分词，构建树节点
        List<String> wordsList = split.apply(words);
        // 匹配字典数据
        Node parent = root;

        for (int i = 0; i < wordsList.size(); i++) {

            String word = wordsList.get(i);

            if (parent.children.containsKey(word)){
                // 存在相同节点则继续向后查找
                parent = parent.children.get(word);
            }else {
                // 如果没有响应子节点就构建新的子节点
                Node node = new Node(word);
                node.isEndWord = (i == (wordsList.size() - 1));
                parent.children.put(word,node);
                parent = node;
            }
        }
    }

    public List<String> search(String words) {

        if (words == null || words.trim().isEmpty()){
            return Collections.emptyList();
        }

        // 分词，构建树节点
        List<String> wordsList = split.apply(words);
        // 匹配字典数据
        Node parent = root;
        int startIndex = 0;

        for (; startIndex < wordsList.size(); startIndex++) {
            String word = wordsList.get(startIndex);
            if (parent.children.containsKey(word)) {
                parent = parent.children.get(word);
            }else {
                break;
            }
        }

        // 如果不最后一个词（也就是中途break了），则返回空列表
        List<String> result = new ArrayList<>();

        if (startIndex == wordsList.size()){
            // 最后一个词则需要恢复句子
            String baseStr = merge.apply(wordsList);
            result = resumeWords(baseStr,parent);
        }

        return result;
    }

    private List<String> resumeWords(String prefix,Node node) {
        List<String> resumeWords = new ArrayList<>();
        if(node.isEndWord){
            resumeWords.add(prefix);

        }

        Collection<Node> values = node.children.values();
        for (Node child : values) {
            String base = merge.apply(List.of(prefix , child.word));
            resumeWords.addAll(resumeWords(base,child));
        }

        return resumeWords;
    }

}
